{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AYV_dMVDxyc2",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "[![Github](https://img.shields.io/github/stars/labmlai/annotated_deep_learning_paper_implementations?style=social)](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/normalization/deep_norm/experiment.ipynb)\n",
    "\n",
    "## DeepNorm\n",
    "\n",
    "This is an experiment training Shakespeare dataset with a deep transformer using [DeepNorm](https://nn.labml.ai/normalization/deep_norm/index.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AahG_i2y5tY9",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Install the packages"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ZCzmCrAIVg0L",
    "outputId": "cf107fb2-4d50-4c67-af34-367624553421",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "source": [
    "!pip install labml-nn --quiet"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SE2VUQ6L5zxI",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "0hJXx_g0wS2C",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from labml import experiment\n",
    "from labml.configs import option\n",
    "from labml_nn.normalization.deep_norm.experiment import Configs"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Lpggo0wM6qb-",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Create an experiment"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "bFcr9k-l4cAg",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "source": [
    "experiment.create(name=\"deep_norm\", writers={'screen'})"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-OnHLi626tJt",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Configurations"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "Piz0c5f44hRo",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "source": [
    "conf = Configs()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wwMzCqpD6vkL",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Set experiment configurations and assign a configurations dictionary to override configurations"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17
    },
    "id": "e6hmQhTw4nks",
    "outputId": "29634715-42f4-4405-fb11-fc9522608627",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "source": [
    "experiment.configs(conf, {\n",
    "    # Use character level tokenizer\n",
    "    'tokenizer': 'character',\n",
    "    # Prompt separator is blank\n",
    "    'prompt_separator': '',\n",
    "    # Starting prompt for sampling\n",
    "    'prompt': 'It is ',\n",
    "    # Use Tiny Shakespeare dataset\n",
    "    'text': 'tiny_shakespeare',\n",
    "\n",
    "    # Use a context size of $256$\n",
    "    'seq_len': 256,\n",
    "    # Train for 32 epochs\n",
    "    'epochs': 32,\n",
    "    # Batch size $16$\n",
    "    'batch_size': 16,\n",
    "    # Switch between training and validation for $10$ times per epoch\n",
    "    'inner_iterations': 10,\n",
    "\n",
    "    # Adam optimizer with no warmup\n",
    "    'optimizer.optimizer': 'Adam',\n",
    "    'optimizer.learning_rate': 3e-4,\n",
    "})"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EvI7MtgJ61w5",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Set PyTorch models for loading and saving"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 255
    },
    "id": "GDlt7dp-5ALt",
    "outputId": "e7548e8f-c541-4618-dc5a-1597cae42003",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "source": [
    "experiment.add_pytorch_models({'model': conf.model})"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KJZRf8527GxL",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Start the experiment and run the training loop."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "aIAWo7Fw5DR8",
    "outputId": "db979785-bfe3-4eda-d3eb-8ccbe61053e5",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "source": [
    "# Start the experiment\n",
    "with experiment.start():\n",
    "    conf.run()"
   ],
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "DeepNorm",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
